# @Time    : 2017/11/10 下午2:59
# @Author  : JCD
# @File    : webdriver_basic.py
from selenium import webdriver
from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
import re
from ratel.common.dict import header_default
from http.cookiejar import CookieJar, Cookie
import time


class Driver:

    def __init__(self, driver_type, executable_path, headers=None, cookies=None, proxy=None, loadimgs=True):

        """
        初始化webdriver
        :param driver_type: 浏览器类型(仅支持'chrome'和'phantomjs')
        :param executable_path: 驱动路径
        :param headers: 请求头(dict类型)
        :param cookies: cookie(list[dict]类型)
        :param proxy: 代理ip(严格遵守格式:xxx.xxx.xxx.xxx:xx)
        :param loadimgs: 是否加载图像(默认加载, 推荐phantomjs禁用图片)
        """
        if executable_path is None or "" == executable_path:
            raise ValueError('executable_path can`t be None!')
        elif not isinstance(executable_path, str):
            raise ValueError('executable_path must be a string!')

        if not(headers is None) and not isinstance(headers, dict):
            raise ValueError('headers must be a dict!')

        if not(cookies is None) and not isinstance(cookies, list):
            raise ValueError('cookies must be a list!')

        if not(proxy is None):
            if not isinstance(proxy, str):
                raise ValueError('proxy must be a string!')
            if "" == proxy:
                raise ValueError('proxy can`t be None!')
            proxy = proxy.replace('localhost', '127.0.0.1')
            if not self.__check_ip(proxy):
                raise ValueError("proxy must be a 'xxx.xxx.xxx.xxx:xx'")
        # 设置允许访问https
        service_args = ['--ssl-protocol=any', '--ignore-ssl-errors=true']
        if not(proxy is None):
            service_args.append('--proxy=' + str(proxy))
            service_args.append('--proxy-type=http')

        if driver_type is None or "" == driver_type:
            raise ValueError('driver_type can`t be None!')
        elif not isinstance(driver_type, str):
            raise ValueError('driver_type must be a string!')
        elif 'CHROME' == driver_type.upper():
            options = webdriver.ChromeOptions()
            # 设置chrome禁用图像
            if not loadimgs:
                prefs = {"profile.managed_default_content_settings.images": 2}
                options.add_experimental_option("prefs", prefs)
            # 设置默认UA
            options.add_argument('--user-agent=' + header_default['Chrome']['User-Agent'])
            # 设置请求头(chrome中只能设置UA)
            if not(headers is None):
                for key, value in headers.items():
                    if 'USER-AGENT' == key.upper():
                        options.add_argumnt('--user-agent=' + value)

            self.__driver = webdriver.Chrome(executable_path=executable_path, chrome_options=options, service_args=service_args)
        elif 'PHANTOMJS' == driver_type.upper():
            des = DesiredCapabilities.CHROME.copy()
            # 设置phantomjs禁用图像
            if not loadimgs:
                des["phantomjs.page.settings.loadImages"] = False
            # 设置默认UA
            des['phantomjs.page.settings.userAgent'] = header_default['Chrome']['User-Agent']
            # 设置请求头
            if not(headers is None):
                for key, value in headers.items():
                    des['phantomjs.page.customHeaders.{}'.format(key)] = value

            self.__driver = webdriver.PhantomJS(executable_path=executable_path, desired_capabilities=des, service_args=service_args)
        else:
            raise ValueError("driver_type only for 'chrome' or 'phantomjs'!")

        if not(cookies is None):
            for cookie in cookies:
                # fix the problem-> "errorMessage":"Unable to set Cookie"
                for k in ('name', 'value', 'domain', 'path', 'expiry'):
                    if k not in list(cookie.keys()):
                        if k == 'expiry':
                            cookie[k] = int(time.time()) + 9999  # 时间戳 秒
                self.__driver.add_cookie({k: cookie[k] for k in ('name', 'value', 'domain', 'path', 'expiry') if k in cookie})

    @property
    def get_driver(self):
        return self.__driver

    def set_driver(self, driver):
        """
        设置webdriver
        :param driver: webdriver实例
        :return:
        """
        self.__driver = driver

    @staticmethod
    def get_cookie2cookiejar(driver):
        cookie_jar = CookieJar()
        for cookie in driver.get_cookies():
            cookie_temp = Cookie(version=0, name=cookie['name'], value=cookie['value'],
                                 port=None, port_specified=False,
                                 domain=cookie['domain'], domain_specified=False, domain_initial_dot=False,
                                 path='', path_specified=False,
                                 secure=False,
                                 expires=int(time.time()) + 9999,
                                 discard=False,
                                 comment=None,
                                 comment_url=None,
                                 rest={},
                                 rfc2109=False)
            cookie_jar.set_cookie(cookie_temp)
        return cookie_jar

    @staticmethod
    def __check_ip(ip):
        """
        校验ip与端口格式
        :param ip:
        :return:
        """
        pattern = '^(25[0-5]|2[0-4][0-9]|[0-1]{1}[0-9]{2}|[1-9]{1}[0-9]{1}|[1-9])\.(25[0-5]|2[0-4][0-9]|' \
                  '[0-1]{1}[0-9]{2}|[1-9]{1}[0-9]{1}|[1-9]|0)\.(25[0-5]|2[0-4][0-9]|[0-1]{1}[0-9]{2}|' \
                  '[1-9]{1}[0-9]{1}|[1-9]|0)\.(25[0-5]|2[0-4][0-9]|[0-1]{1}[0-9]{2}|[1-9]{1}[0-9]{1}|[0-9])' \
                  ':(([0-9][0-9]{0,3})|([1-5]\d{4})|(6[1-4]\d{3})|(65[1-4]{2})|(655[1-2]\d)|(6553[1-5]))$'
        return True if re.match(pattern, ip) else False